#pragma once
#include <stl/trie_node.h>
#include <string>
#include <type_traits>

using std::string;

namespace stl 
{
	template<typename T>
	class Trie
	{
		typedef TrieNode<T> Node;
	public:
		static_assert(std::is_same<T, uint8_t>::value 
				|| std::is_same<T, uint16_t>::value 
				|| std::is_same<T, uint32_t>::value 
				|| std::is_same<T, uint64_t>::value 
				|| std::is_same<T, string>::value,
				"trie suppose only uint8, uint16, uint32, uint64, string type");

	public:
		Trie() : m_root(nullptr), m_size(0) {
			m_root = new Node();
		}

		~Trie() {
			if (m_root != nullptr) {
				delete m_root;
				m_root = nullptr;
			}
			m_size = 0;
		}

		// 字符串类型的插入
		template<typename U = T>
		typename std::enable_if<std::is_same<U, string>::value, void>::type
		insert(const U &key) {
			insert(key.data(), key.size());
		} 

		// 字符串字面量的插入
		void insert(const char* key) {
			insert(key, strlen(key));
		}

		// 数值类型的插入
		template<typename U = T>
		typename std::enable_if<!std::is_same<U, string>::value, void>::type
		insert(const U &key) {
			string s = to_string(key);
			insert(s.data(), s.size());
		}

		int size() const {
			return m_size;
		}
 
		// 查找字符串前缀是否存在
		template<typename U = T>
		typename std::enable_if<std::is_same<U, string>::value, bool>::type
		prefix(const U &key) {
			auto node = searchNode(key.data(), key.size());
			return node != nullptr;
		}

		// 查找数字类型前缀是否存在
		template<typename U = T>
		typename std::enable_if<!std::is_same<U, string>::value, bool>::type
		prefix(const U &key) {
			string s = to_string(key);
			auto node = searchNode(s.data(), s.size());
			return node != nullptr;
		}

		// 查找某个单词是否存在
		template<typename U = T>
		typename std::enable_if<std::is_same<U, string>::value, bool>::type
		search(const U &key) {
			auto node = searchNode(key.data(), key.size());
			if (!node) return false;
			return node->m_flag;
		}

		// 字符串字面量查找
		bool search(const char* key) {
			auto node = searchNode(key, strlen(key));
			if (!node) return false;
			return node->m_flag;
		}

		// 查找某个数字是否存在
		template<typename U = T>
		typename std::enable_if<!std::is_same<U, string>::value, bool>::type
		search(const U &key) {
			string s = std::to_string(key);
			auto node = searchNode(s.data(), s.size());
			if (!node) return false;
			return node->m_flag;
		}

	protected:
		void insert(const char * data, int size) {			
			Node* node = m_root;
			bool isNewWord = false;

			for (int i = 0; i < size; i++) {
				char c = data[i];
				int idx = 0;
				bool found = false;

				while (idx < node->m_size) {
					if (node->m_children[idx]->m_data == c) {
						found = true;
						break;
					} else if (node->m_children[idx]->m_data > c) {
						break;
					}
					idx++;
				}
				
				if (found) {
					node = node->m_children[idx];
				} else {
					Node* newNode = new Node();
					newNode->m_data = c;
					node->insert(idx, newNode);
					node = newNode;
					isNewWord = true;
				}
			}
			
			// 到达单词结尾，如果不是单词结尾节点，标记为结尾
			if (!node->m_flag) {
				node->m_flag = true;
				m_size++; // 只有当这是一个新单词时才增加计数
			}
		}

		Node *searchNode(const char *data, int size) {
			Node* node = m_root;

			for (int i = 0; i < size; i++) {
				char c = data[i];
				int idx = 0;
				bool found = false;

				while (idx < node->m_size) {
					if (node->m_children[idx]->m_data == c) {
						found = true;
						break;
					} else if (node->m_children[idx]->m_data > c) {
						break;
					}
					idx++;
				}
				if (!found) return nullptr;
				node = node->m_children[idx];
			}
			return node;
		}

	private:
		Node *m_root;
		int m_size;
	};
}